clc;
clear;

tic;
startTime = cputime;
pde = process_pde_2d(1);
vector_h = [1/8 1/16 1/32];
iter_L = 8;

u_Linf = zeros(size(vector_h,2),1);
u_L2 = zeros(size(vector_h,2),1);
u_H1 = zeros(size(vector_h,2),1);
p_Linf = zeros(size(vector_h,2),1);
p_L2 = zeros(size(vector_h,2),1);
p_H1 = zeros(size(vector_h,2),1);

for i = 1:size(vector_h,2)
    h = vector_h(i);
    mesh = process_mesh_h_2d(pde, h, "P2");

    A1 = assemble_stiffness_matrix_2d(mesh, "P2", "dx", "P2", "dx", pde.nu);
    A2 = assemble_stiffness_matrix_2d(mesh, "P2", "dy", "P2", "dy", pde.nu);
    A3 = assemble_stiffness_matrix_2d(mesh, "P2", "dx", "P2", "dy", pde.nu);
    A5 = assemble_stiffness_matrix_2d(mesh, "P1", "x", "P2", "dx", -1);
    A6 = assemble_stiffness_matrix_2d(mesh, "P1", "x", "P2", "dy", -1);
    % A0 = assemble_stiffness_matrix_2d(mesh, "P1", "x", "P1", "x", -1e-6);
    % A0 = 1e-8*speye(mesh.Np);
    % A0 = zeros(mesh.Np);
    A01 = sparse(mesh.Np, mesh.Np);
    A02 = sparse(mesh.Nb, mesh.Np);
    A = [2*A1+A2, A3, A5;
        A3', 2*A2+A1, A6;
        A5', A6', A01;];
    [b1, b2] = assemble_load_vector_2d(mesh, pde);
    b0 = zeros(mesh.Np, 1);
    b = [b1; b2; b0];

    U_old = zeros(2*mesh.Nb+mesh.Np, 1);
    for l = 1:iter_L
        AN1 = assemble_nonlinear_matrix_2d(mesh, "u1", "dx", "P2", "x", U_old);
        AN2 = assemble_nonlinear_matrix_2d(mesh, "u1", "x", "P2", "dx", U_old);
        AN3 = assemble_nonlinear_matrix_2d(mesh, "u2", "x", "P2", "dy", U_old);
        AN4 = assemble_nonlinear_matrix_2d(mesh, "u1", "dy", "P2", "x", U_old);
        AN5 = assemble_nonlinear_matrix_2d(mesh, "u2", "dx", "P2", "x", U_old);
        AN6 = assemble_nonlinear_matrix_2d(mesh, "u2", "dy", "P2", "x", U_old);
        AN = [AN1+AN2+AN3, AN4, A02;
            AN5, AN6+AN2+AN3, A02;
            A02', A02', A01];
        bN1 = assemble_nonlinear_vector_2d(mesh, "u1", "u1", "dx", U_old);
        bN2 = assemble_nonlinear_vector_2d(mesh, "u2", "u1", "dy", U_old);
        bN3 = assemble_nonlinear_vector_2d(mesh, "u1", "u2", "dx", U_old);
        bN4 = assemble_nonlinear_vector_2d(mesh, "u2", "u2", "dy", U_old);
        bN = [bN1+bN2;
            bN3+bN4;
            b0];

        Al = A + AN;
        bl = b + bN;
        [Al, bl] = process_boundary_2d(Al, bl, mesh, pde);
        Al(2*mesh.Nb+1,:) = 0;
        Al(2*mesh.Nb+1,2*mesh.Nb+1) = 1;
        bl(2*mesh.Nb+1,1) = pde.p(pde.domain(1), pde.domain(3));
        % A(end,:) = 0;
        % A(end,end) = 1;
        % b(end,1) = pde.p(pde.domain(2), pde.domain(4));
        % cond(A)
        U = Al\bl;
        % U = gmres(A, b);
        % U = minres(A, b);
        % U = lsqminnorm(A, b);
        % U = lsqr(A, b);

        U_old = U;
    end


    u_Linf(i,1) = process_error_u_2d(mesh, pde, U, "Linf");
    u_L2(i,1) = process_error_u_2d(mesh, pde, U, "L2");
    u_H1(i,1) = process_error_u_2d(mesh, pde, U, "H1");
    p_Linf(i,1) = process_error_p_2d(mesh, pde, U, "Linf");
    p_L2(i,1) = process_error_p_2d(mesh, pde, U, "L2");
    p_H1(i,1) = process_error_p_2d(mesh, pde, U, "H1");

    fprintf("完成\t%d/%d\n", i, size(vector_h,2));
end

process_convergence_h_2d(u_Linf, vector_h);
process_convergence_h_2d(u_L2, vector_h);
process_convergence_h_2d(u_H1, vector_h);
process_convergence_h_2d(p_Linf, vector_h);
process_convergence_h_2d(p_L2, vector_h);
process_convergence_h_2d(p_H1, vector_h);

elapsedTOCTime = toc;
elapsedCPUTime = cputime - startTime;
disp(["TOC time(s)", num2str(elapsedTOCTime)]);
disp(["CPU time(s)", num2str(elapsedCPUTime)]);

%% assemble_stiffness_matrix_2d
function matrix_A = assemble_stiffness_matrix_2d(mesh, B1_type, D1_type, B2_type, D2_type, cor)
basis_map = containers.Map(...
    {'P1', 'P1b', 'P2'},...
    {struct('Nb', mesh.Np, 'Nlb',3),...
    struct('Nb', mesh.Nb, 'Nlb',4),...
    struct('Nb', mesh.Nb, 'Nlb',6)});
B1 = basis_map(B1_type);
Nb1 = B1.Nb();
Nlb1 = B1.Nlb();
B2 = basis_map(B2_type);
Nb2 = B2.Nb();
Nlb2 = B2.Nlb();
rows = zeros(2*Nlb1*Nlb2,1);
cols = rows;
values = rows;
flag = 1;
for n = 1:mesh.Ne
    E = mesh.node(mesh.elem(n,:),:);
    V = E(1:3,:);
    [gauss_points_loc, gauss_weights_loc] = gauss_lengdre_loc(V, "2d");
    phi1 = basis_function_loc(V, gauss_points_loc(:,1), gauss_points_loc(:,2), B1_type, D1_type);
    phi2 = basis_function_loc(V, gauss_points_loc(:,1), gauss_points_loc(:,2), B2_type, D2_type);
    for i = 1:Nlb2
        for j = 1:Nlb1
            f = cor * phi1{j} .* phi2{i};
            r = gauss_weights_loc * f;
            rows(flag) = mesh.elem(n,i);
            cols(flag) = mesh.elem(n,j);
            values(flag) = r;
            flag = flag + 1;
        end
    end
end
matrix_A = sparse(rows, cols, values, Nb2, Nb1);
end

%% assemble_load_vector_2d
function [vector_b1, vector_b2] = assemble_load_vector_2d(mesh, pde)
vector_b1 = zeros(mesh.Nb,1);
vector_b2 = vector_b1;
for n = 1:mesh.Ne
    E = mesh.node(mesh.elem(n,:),:);
    V = E(1:3,:);
    [gauss_points_loc, gauss_weights_loc] = gauss_lengdre_loc(V, "2d");
    phi = basis_function_loc(V, gauss_points_loc(:,1), gauss_points_loc(:,2), mesh.basis_type);
    ff1 = pde.f1(gauss_points_loc(:,1), gauss_points_loc(:,2));
    ff2 = pde.f2(gauss_points_loc(:,1), gauss_points_loc(:,2));
    for i = 1:mesh.Nlb
        f1 = ff1 .* phi{i};
        f2 = ff2 .* phi{i};
        r1 = gauss_weights_loc * f1;
        r2 = gauss_weights_loc * f2;
        vector_b1(mesh.elem(n,i),1) = vector_b1(mesh.elem(n,i),1) + r1;
        vector_b2(mesh.elem(n,i),1) = vector_b2(mesh.elem(n,i),1) + r2;
    end
end
end

%% assemble_nonlinear_matrix_2d
function matrix_A = assemble_nonlinear_matrix_2d(mesh, u_type, du_type, B1_type, D1_type, cor)
basis_map = containers.Map(...
    {'P1', 'P1b', 'P2'},...
    {struct('Nb', mesh.Np, 'Nlb',3),...
    struct('Nb', mesh.Nb, 'Nlb',4),...
    struct('Nb', mesh.Nb, 'Nlb',6)});
B1 = basis_map(B1_type);
Nb1 = B1.Nb();
Nlb1 = B1.Nlb();
Nb2 = mesh.Nb;
Nlb2 = 6;
rows = zeros(2*Nlb1*Nlb2,1);
cols = rows;
values = rows;
if u_type == "u1"
    shift = 0;
elseif u_type == "u2"
    shift = mesh.Nb;
end
flag = 1;
for n = 1:mesh.Ne
    E = mesh.node(mesh.elem(n,:),:);
    V = E(1:3,:);
    [gauss_points_loc, gauss_weights_loc] = gauss_lengdre_loc(V, "2d");
    phiu = basis_function_loc(V, gauss_points_loc(:,1), gauss_points_loc(:,2), "P1", du_type);
    phi1 = basis_function_loc(V, gauss_points_loc(:,1), gauss_points_loc(:,2), B1_type, D1_type);
    phi2 = basis_function_loc(V, gauss_points_loc(:,1), gauss_points_loc(:,2), "P2");
    du = zeros(9,1);
    for i = 1:3
        du = du + cor(mesh.elem(n,i)+shift) .* phiu{i};
    end
    for i = 1:Nlb2
        for j = 1:Nlb1
            f = du .* phi1{j} .* phi2{i};
            r = gauss_weights_loc * f;
            rows(flag) = mesh.elem(n,i);
            cols(flag) = mesh.elem(n,j);
            values(flag) = r;
            flag = flag + 1;
        end
    end
end
matrix_A = sparse(rows, cols, values, Nb2, Nb1);
end

%% assemble_nonlinear_vector_2d
function vector_b = assemble_nonlinear_vector_2d(mesh, u1_type, u2_type, du2_type, cor)
u_map = containers.Map(...
    {'u1', 'u2'},...
    {struct('shift', 0),...
    struct('shift', mesh.Nb)});
U1 = u_map(u1_type);
shift1 = U1.shift;
U2 = u_map(u2_type);
shift2 = U2.shift;
vector_b = zeros(mesh.Nb,1);
for n = 1:mesh.Ne
    E = mesh.node(mesh.elem(n,:),:);
    V = E(1:3,:);
    [gauss_points_loc, gauss_weights_loc] = gauss_lengdre_loc(V, "2d");
    phi1 = basis_function_loc(V, gauss_points_loc(:,1), gauss_points_loc(:,2), "P1");
    phi2 = basis_function_loc(V, gauss_points_loc(:,1), gauss_points_loc(:,2), "P1", du2_type);
    phi3 = basis_function_loc(V, gauss_points_loc(:,1), gauss_points_loc(:,2), "P2");
    item1 = zeros(9,1);
    item2 = zeros(9,1);
    for i = 1:3
        item1 = item1 + cor(mesh.elem(n,i)+shift1) .* phi1{i};
        item2 = item2 + cor(mesh.elem(n,i)+shift2) .* phi2{i};
    end
    for j = 1:mesh.Nlb
        f =  item1 .* item2 .* phi3{j};
        r = gauss_weights_loc * f;
        vector_b(mesh.elem(n,j),1) = vector_b(mesh.elem(n,j),1) + r;
    end
end
end